from flask import Flask, render_template, Response,request
import cv2
import ffmpeg
import numpy as np

class VideoCamera(object):
    def __init__(self,name):
        self.camera = 'rtsp://36.133.66.40:8554/'+name

        self.probe = ffmpeg.probe(self.camera)
        self.video_stream = next((stream for stream in self.probe['streams'] if stream['codec_type'] == 'video'), None)
        self.width = int(self.video_stream['width'])
        self.height = int(self.video_stream['height'])

        self.out = (
            ffmpeg
                .input(self.camera, rtsp_transport='tcp')
                .output('pipe:', format='rawvideo', pix_fmt='bgr24', loglevel="quiet", r=25)
                .run_async(pipe_stdout=True)
        )
        self.cnt_empty = 0

    def __del__(self):
        pass

    def get_frame(self):
        in_bytes = self.out.stdout.read(self.height * self.width * 3)
        if not in_bytes:
            self.cnt_empty += 1
            if self.cnt_empty > 10:
                return 0
        cnt_empty = 0
        frame = np.frombuffer(in_bytes, dtype=np.uint8).reshape(self.height, self.width, 3)
        if len(frame) > 0:
            ret, jpeg = cv2.imencode('.jpg', frame)
            return jpeg.tobytes()
        else:
            return True


app = Flask(__name__)

@app.route('/')
def index():
    # jinja2模板，具体格式保存在index.html文件中
    return render_template('display.html')


def gen(camera):
    while True:
        frame = camera.get_frame()
        # 使用generator函数输出视频流， 每次请求输出的content类型是image/jpeg
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n\r\n')


@app.route('/video_feed')  # 这个地址返回视频流响应
def video_feed():
    try:
        name = request.args.get("name","")
        return Response(gen(VideoCamera(name)),mimetype='multipart/x-mixed-replace; boundary=frame')
    except:
        pass


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)
